#include "util/arena/arena_wrapper.h"

#include <memory>

#include <bvar/bvar.h>
#include <gflags/gflags.h>

// copy from hawking util to avoid submodule
namespace hawking {
namespace indexlib {

DEFINE_int64(protobuf_arena_start_block_size, 256,
             "Default value is same as google arena default value.");
DEFINE_int64(protobuf_arena_max_block_size, 8192,
             "Default value is same as google arena default value.");
DEFINE_int64(tl_exclusive_arena_initial_block_size, 1024*1024, "");

namespace {

std::once_flag g_arena_monitor_once;

class ArenaMonitor {
public:
    void Hide() {
        lr_.hide();
        count_status_.hide();
    }

    void IncreaseCount() {
        count_.fetch_add(1, std::memory_order_relaxed);
    }

    void DecreaseCount() {
        count_.fetch_sub(1, std::memory_order_relaxed);
    }

    void RecordAllocatedSpace(int64_t space) {
        lr_ << space;
    }

    static ArenaMonitor* Instance() {
        static ArenaMonitor* monitor = nullptr;
        std::call_once(
            g_arena_monitor_once, []() { monitor = new ArenaMonitor(); });
        return monitor;
    }

    static int GetCount(void* monitor) {
        return reinterpret_cast<ArenaMonitor*>(monitor)->count_.load(
            std::memory_order_relaxed);
    }

private:
    ArenaMonitor() : lr_("arena_allocated_space")
                   , count_status_("arena_inflight_count", GetCount, this) {}

    std::atomic<int> count_{0};
    bvar::LatencyRecorder lr_;
    bvar::PassiveStatus<int> count_status_;
};

}

ArenaWrapper::ArenaWrapper(std::shared_ptr<google::protobuf::Arena> arena)
    : arena_(arena) {}

ArenaWrapper::ArenaWrapper() : arena_(CreateArena()) {
}

ArenaWrapper::ArenaWrapper(const google::protobuf::ArenaOptions& options)
    : arena_(new google::protobuf::Arena(options), &DestroyArena) {
    ArenaMonitor::Instance()->IncreaseCount();
}

ArenaWrapper::~ArenaWrapper() = default;

void ArenaWrapper::DestroyArena(google::protobuf::Arena* arena) {
    if (arena != nullptr) {
        ArenaMonitor::Instance()->DecreaseCount();
        ArenaMonitor::Instance()->RecordAllocatedSpace(arena->SpaceAllocated());
        // TODO(caidaojin): optimize if destroy cost is expensive.
        delete arena;
    }
}

std::shared_ptr<google::protobuf::Arena> ArenaWrapper::CreateArena() {
    google::protobuf::ArenaOptions options;
    options.start_block_size = FLAGS_protobuf_arena_start_block_size;
    options.max_block_size = FLAGS_protobuf_arena_max_block_size;
    std::shared_ptr<google::protobuf::Arena> arena(
        new google::protobuf::Arena(options), &DestroyArena);
    ArenaMonitor::Instance()->IncreaseCount();
    return arena;
}

void ArenaWrapper::CloseMonitor() {
    ArenaMonitor::Instance()->Hide();
}


thread_local char* g_tl_exclusive_pb_arena_initial_block = nullptr;
thread_local int g_tl_exclusive_pb_reference = 0;

void CreateThreadLocalExclusiveArena(
    std::unique_ptr<indexlib::ArenaWrapper>* arena, char** arena_init_block) {
    CHECK_EQ(g_tl_exclusive_pb_reference++, 0);
    if (g_tl_exclusive_pb_arena_initial_block == nullptr) {
        g_tl_exclusive_pb_arena_initial_block =
            new char[FLAGS_tl_exclusive_arena_initial_block_size];
    }
    google::protobuf::ArenaOptions options;
    options.initial_block = g_tl_exclusive_pb_arena_initial_block;
    options.initial_block_size = FLAGS_tl_exclusive_arena_initial_block_size;
    *arena_init_block = g_tl_exclusive_pb_arena_initial_block;
    *arena = std::make_unique<indexlib::ArenaWrapper>(options);
}

void DestryThreadLocalExclusiveArena(
    std::unique_ptr<indexlib::ArenaWrapper>& arena,
    char* arena_init_block) {
    if (arena != nullptr ) {
        CHECK_EQ(--g_tl_exclusive_pb_reference, 0);
        CHECK(g_tl_exclusive_pb_arena_initial_block == arena_init_block);
        arena.reset();
    }
}

} // namespace indexlib
} // namespace hawking